import os
import csv
import random
import json

def node_pair(n):
    with open('./ngram_statistics/radius_100/trigram/node_pair.csv', 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        for i in range(400):
            filename = './place/1min_100m/' + str(i) + '.csv'
            if not os.path.exists(filename):
                pass
            else:
                with open(filename, "r") as f:
                    reader = csv.reader(f)
                    for trajectories in reader:
                        if len(trajectories) < n:
                            pass
                        for i in range(len(trajectories)-2):
                            spamwriter.writerow([int(trajectories[i]), int(trajectories[i+1]), int(trajectories[i+2])])
                        spamwriter.writerow([int(trajectories[len(trajectories)-2]), int(trajectories[len(trajectories)-1]), 0])


def node_count():
    nodes = [0]*7050
    filename = './ngram_statistics/radius_100/train4000k.csv'
    filename2 = './ngram_statistics/radius_100/node_count_train4000k.csv'
        
    with open(filename, "r") as f:
        reader = csv.reader(f)
        for trajectories in reader:
            nodes[int(trajectories[0])] += 1

    with open(filename2, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        index = 0
        for item in nodes:
            index += 1
            if item > 0:
                spamwriter.writerow([index, item])

def node_count_2():
    trigram = {}
    filename = './ngram_statistics/radius_100/trigram/train4000k.csv'
    filename2 = './ngram_statistics/radius_100/trigram/node_count_train4000k.json'
        
    with open(filename, "r") as f:
        reader = csv.reader(f)
        for trajectories in reader:
            prev = trajectories[0]
            cur = trajectories[1]
            k = prev+'_'+cur
            if not trigram.__contains__(k):
                trigram[k] = {'CNT':0}
            trigram[k]['CNT'] += 1

    json_data = json.dumps(trigram)
    with open(filename2, 'w') as f3:
        f3.write(json_data)

def divide_dataset():
    filename = './ngram_statistics/radius_100/trigram/node_pair.csv'
    filename2 = './ngram_statistics/radius_100/trigram/train1000k.csv'
    filename3 = './ngram_statistics/radius_100/trigram/train2000k.csv'
    filename4 = './ngram_statistics/radius_100/trigram/train3000k.csv'
    filename5 = './ngram_statistics/radius_100/trigram/train4000k.csv'
    filename6 = './ngram_statistics/radius_100/trigram/train5000k.csv'
    filename7 = './ngram_statistics/radius_100/trigram/test10k.csv'
    filename8 = './ngram_statistics/radius_100/trigram/shuffle_pair.csv'
    
    with open(filename, "r") as f:
        reader = csv.reader(f)
        trajectories = []
        for row in reader:
            trajectories.append(row)
        random.shuffle(trajectories)
    test_set = trajectories[:10000]
    train_set1 = trajectories[10000:1010000]
    train_set2 = trajectories[10000:2010000]
    train_set3 = trajectories[10000:3010000]
    train_set4 = trajectories[10000:4010000]
    train_set5 = trajectories[10000:5010000]
    
    with open(filename2, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        spamwriter.writerows(train_set1)

    with open(filename3, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        spamwriter.writerows(train_set2)

    with open(filename4, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        spamwriter.writerows(train_set3)

    with open(filename5, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        spamwriter.writerows(train_set4)

    with open(filename6, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        spamwriter.writerows(train_set5)

    with open(filename7, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        spamwriter.writerows(test_set)

    with open(filename8, 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile)
        spamwriter.writerows(trajectories)


def MMC2():
    prob_matrix = {}
    node_index = []
    with open("./ngram_statistics/radius_100/node_count_train4000k.csv", "r") as f1:
        reader = csv.reader(f1)
        for item in reader:
            # nodeID = int(item[0])
            nodeCNT = int(item[1])
            prob_matrix[item[0]] = {}
            prob_matrix[item[0]]['CNT'] = nodeCNT
    
    with open("./ngram_statistics/radius_100/train4000k.csv", "r") as f2:
        reader = csv.reader(f2)
        for row in reader:
            if not prob_matrix.__contains__(row[0]):
                continue
            if prob_matrix[row[0]].__contains__(row[1]):
                prob_matrix[row[0]][row[1]] += 1
            else:
                prob_matrix[row[0]][row[1]] = 1
    
    json_data = json.dumps(prob_matrix)
    
    with open('./ngram_statistics/radius_100/KMM2_4000k.json', 'w') as f3:
        f3.write(json_data)


def MMC3():
    with open('./ngram_statistics/radius_100/trigram/node_count_train4000k.json') as bf:
        prob_matrix = json.load(bf)

    node_index = []

    with open("./ngram_statistics/radius_100/trigram/train4000k.csv", "r") as f2:
        reader = csv.reader(f2)
        for row in reader:
            prev = row[0]
            cur = row[1]
            k = prev+'_'+cur
            if not prob_matrix.__contains__(k):
                continue
            if prob_matrix[k].__contains__(row[2]):
                prob_matrix[k][row[2]] += 1
            else:
                prob_matrix[k][row[2]] = 1
    
    json_data = json.dumps(prob_matrix)
    
    with open('./ngram_statistics/radius_100/trigram/KMM3_4000k.json', 'w') as f3:
        f3.write(json_data)


def pred_loc_count():
    with open('./ngram_statistics/radius_100/trigram/KMM3_4000k.json') as bf:
        prob_matrix = json.load(bf)

    loc_count = {}
    for key in prob_matrix:
        prob_list = sorted(prob_matix[key].items(), key = lambda item:item[1], reverse=True)
        prob_cnt = [[0,0],[0,0],[0,0],[0,0],[0,0]]
        pred_cnt = 0
        for i in range(1,len(prob_list)):
            pred_cnt += prob_list[i][1]
            prob = pred_cnt/prob_list[0][1]
            if prob >= 0.6 and prob_cnt[0][1] == 0:
                prob_cnt[0][0] = i
                prob_cnt[0][1] = 1
            if prob >= 0.7 and prob_cnt[1][1] == 0:
                prob_cnt[1][0] = i
                prob_cnt[1][1] = 1
            if prob >= 0.8 and prob_cnt[2][1] == 0:
                prob_cnt[2][0] = i
                prob_cnt[2][1] = 1
            if prob >= 0.9 and prob_cnt[3][1] == 0:
                prob_cnt[3][0] = i
                prob_cnt[3][1] = 1
                break
        prob_cnt[4][0] = len(prob_list)-1
        prob_cnt[4][1] = 1
        loc_count[key] = [prob_cnt[0][0],prob_cnt[1][0],prob_cnt[2][0],prob_cnt[3][0],prob_cnt[4][0]]

    json_data = json.dumps(loc_count)
    
    with open('./ngram_statistics/radius_100/trigram/KMM3_4000k_prob_cnt.json', 'w') as f3:
        f3.write(json_data)


def pred_loc_stat():
    with open('./ngram_statistics/radius_100/trigram/KMM3_4000k_prob_cnt.json') as bf:
        loc_count = json.load(bf)

    loc_stat = {'0':{}, '1':{}, '2':{}, '3':{}, '4':{}}
    for key in loc_count:
        item = loc_count[key]
        for i in range(len(item)):
            val = item[i]
            if not loc_stat[str(i)].__contains__(val):
                loc_stat[str(i)][val] = 0
            loc_stat[str(i)][val] += 1

    json_data = json.dumps(loc_stat)
    
    with open('./ngram_statistics/radius_100/trigram/KMM3_4000k_prob_stat.json', 'w') as f3:
        f3.write(json_data)

if __name__ == '__main__':
    pred_loc_stat()